import pandas as pd
from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error

import input, preprocess

def add_new_country_to_medal_counts(dataframes, noc_code, start_time):
    """
    在medal_counts中添加一个新国家的数据。
    :param dataframes: 包含所有数据框的字典
    :param noc_code: 新国家的NOC代码
    """
    if "medal_counts" in dataframes and dataframes["medal_counts"] is not None:
        medal_counts = dataframes["medal_counts"].copy()
        
        # 生成新国家的数据
        years = list(range(start_time, 2025, 4))  # 从start_time到2024，每4年一次
        new_data = {
            "Rank": [999] * len(years),
            "NOC": noc_code,
            "Gold": [0] * len(years),
            "Silver": [0] * len(years),
            "Bronze": [0] * len(years),
            "Total": [0] * len(years),
            "Year": years,
        }

        new_df = pd.DataFrame(new_data)
        
        # 合并新数据到现有的medal_counts数据框中
        medal_counts = pd.concat([medal_counts, new_df], ignore_index=True)
        # medal_counts.to_csv("./medal_counts.csv", index=False, encoding="utf-8")
        # print("预测结果已保存到 medal_counts.csv 文件中。")
        
        # 更新数据框
        dataframes["medal_counts"] = medal_counts
    else:
        print("Error: 'medal_counts' DataFrame is not available.")


# 模型训练和预测
def first_medal_train_and_predict(features, year_to_predict):
    # 特征选择
    features_to_use = ["Year", "Is_Host", "Last_Total_Medals", "Total_Programs"]
    # 筛选历史数据
    historical_data = features[features["Year"] < year_to_predict]
    # 特征和目标变量
    X = historical_data[features_to_use]
    y = historical_data["Athlete_Win_Ratio"]
    # 确保特征数据为数值类型
    X = X.apply(pd.to_numeric, errors='coerce').fillna(0)
    y = pd.to_numeric(y, errors='coerce').fillna(0)
    # 数据分割
    X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=42)
    # 模型训练
    model = RandomForestRegressor(n_estimators=100, random_state=42)
    model.fit(X_train, y_train)
    # 模型评估
    y_pred = model.predict(X_val)
    mse = mean_squared_error(y_val, y_pred)
    print(f"\nMean Squared Error (Validation Set): \033[91m{mse}\033[0m")
    # 获取本次奥运会的特征数据
    latest_data = features[features["Year"] == (year_to_predict - 4)].copy()
    latest_data["Year"] = year_to_predict
    latest_data["Is_Host"] = (latest_data["Country"] == host_country).astype(int)
    latest_data["Last_Total_Medals"] = latest_data["Total"]
    latest_data["Total_Programs"] = num_of_programs
    # 用本次奥运会的特征数据作为x来预测 Athlete_Win_Ratio
    future_X = latest_data[features_to_use]
    future_X = future_X.apply(pd.to_numeric, errors='coerce').fillna(0)
    predicted_athlete_win_ratio = model.predict(future_X)
    # 将预测值保存到 latest_data 中
    latest_data["Predicted_Athlete_Win_Ratio"] = predicted_athlete_win_ratio
    # 确保预测值在 0 到 1 之间
    latest_data["Predicted_Athlete_Win_Ratio"] = latest_data["Predicted_Athlete_Win_Ratio"].clip(0, 1)
    # 如果预测值过低，可以尝试调整
    latest_data["Predicted_Athlete_Win_Ratio"] = latest_data["Predicted_Athlete_Win_Ratio"].apply(lambda x: max(x, 0.1))
    
    return latest_data


if __name__ == "__main__":
    file_path = "./data/"  # 替换为你的文件路径
    dataframes = input.load_data(file_path)
    # input.print_dataframes(dataframes) # 打印输入的数据表

    # 添加国家“a_new_country”
    start_time_1 = 2000
    noc_code_1 = "a_new_country"
    add_new_country_to_medal_counts(dataframes, noc_code_1, start_time_1)

    # # 添加2024年国家“Belarus”
    # start_time_2 = 2024
    # noc_code_2 = "Belarus"
    # add_new_country_to_medal_counts(dataframes, noc_code_2, start_time_2)

    # 预处理，提取特征
    features = preprocess.preprocess_data(dataframes)

    # 输入需要预测的年份，和主办国家
    year_to_predict = 2028
    host_country = "United States"
    num_of_programs = 350  # 请根据实际情况调整
    
    # 训练模型并预测
    latest_data = first_medal_train_and_predict(features, year_to_predict)
    # 输出预测结果
    print("\nPredicted Athlete Win Ratio for", year_to_predict, "Olympics:")
    print(latest_data[["Country", "Predicted_Athlete_Win_Ratio"]])
    # 保存预测结果
    latest_data.to_csv("./predictions.csv", index=False, encoding="utf-8")
    print("预测结果已保存到 predictions.csv 文件中。")